{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 导入相关包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "from datasets import DatasetDict\n",
    "from transformers import AutoTokenizer,AutoModelForTokenClassification,DataCollatorForTokenClassification,TrainingArguments,Trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 获取实体标注类型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ner_datasets = DatasetDict.load_from_disk(\"/data/datasets/ner_data\")\n",
    "label_list = ner_datasets['train'].features[\"ner_tags\"].feature.names\n",
    "label_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 数据处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"/data/models/merged_model_2_02\")\n",
    "#  is_split_into_words 决定了词和句子的处理方式\n",
    "res = tokenizer([\"10.10.112.59\"]) # 7个词 ，7个标注 ，wordids 有 7 个。\n",
    "print(res.input_ids)\n",
    "print(res.word_ids())\n",
    "\n",
    "res = tokenizer([\"interesting word\"]) # 2个词，2 个标注 ，wordids 有 5 个。\n",
    "print(\"输入id序列\",res.input_ids)\n",
    "print(\"词的id\",res.word_ids())\n",
    "\n",
    "res = tokenizer(ner_datasets['train'][0:2][\"tokens\"],is_split_into_words=True,padding='max_length',max_length=64,truncation=True,return_offsets_mapping=True,return_tensors='pt')\n",
    "ner_datasets['train'][0:1][\"ner_tags\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model =  AutoModelForTokenClassification.from_pretrained(\"/data/models/merged_model_2_02\",num_labels=len(label_list))\n",
    "model.to(\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 模型评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"https://nktgfec7cyfqw8k.66776699.com/ 是一个错误的网站，硬盘应该访问 www.baidu.com。\"\n",
    "text = \"中国乒乓球协会将继续致力于中国乒乓球事业健康发展，进一步弘扬新时代中华体育精神的新风貌，助力体育强国建设，全力为祖国和人民赢得更大的荣誉。\"\n",
    "tokenizer_output =  tokenizer([[text]],max_length=128, truncation=True, is_split_into_words=True)\n",
    "output = model(\n",
    "        input_ids=torch.tensor(tokenizer_output[\"input_ids\"]).to(\"cuda\"),\n",
    "        attention_mask=torch.tensor(tokenizer_output[\"attention_mask\"]).to(\"cuda\"),\n",
    "        token_type_ids=torch.tensor(tokenizer_output[\"token_type_ids\"]).to(\"cuda\"),\n",
    ")\n",
    "logits = output.logits[0].detach().cpu().numpy()\n",
    "# print(logits)\n",
    "predictions =  np.argmax(logits,axis=-1)\n",
    "for i in tokenizer_output[\"input_ids\"]:\n",
    "    print(tokenizer.decode(i))\n",
    "print(predictions)\n",
    "# ['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_entities(predictions,label_list,tokenizer,tokenizer_output):\n",
    "    entities = []\n",
    "    entity_types = []\n",
    "\n",
    "    for i in range(len(predictions)):\n",
    "        if predictions[i] != 0:\n",
    "            if label_list[predictions[i]].startswith(\"B-\"):\n",
    "                start = []\n",
    "                entities.append(start)\n",
    "                entity_types.append(label_list[predictions[i]][2:])\n",
    "                start.append(tokenizer.decode(tokenizer_output[\"input_ids\"][0][i]))\n",
    "            else:\n",
    "                start.append(tokenizer.decode(tokenizer_output[\"input_ids\"][0][i]))\n",
    "    return [\"\".join(i) for i in entities],entity_types\n",
    "\n",
    "entities,entity_types = get_entities(predictions,label_list,tokenizer,tokenizer_output)\n",
    "print(entities,entity_types)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ner_datasets = DatasetDict.load_from_disk(\"/data/datasets/ner_data\")\n",
    "label_list = ner_datasets['train'].features[\"ner_tags\"].feature.names\n",
    "label_list"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
